
"""
构建数据分析流程图
"""
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from core.state import AnalyzeState
from core.nodes import (
    userinfo_node, schemas_node, rewrite_query_node,
    wyx_rewrite_query_node,
    biw_query_node, dataset_filter_node,
)


def _build_base_graph():
    """Build and return the base state graph with all nodes and edges."""
    builder = StateGraph(AnalyzeState)

    builder.add_edge(START, "userinfo")
    builder.add_node("userinfo", userinfo_node)
    builder.add_node("schemas", schemas_node)
    builder.add_node("rewrite_query", rewrite_query_node)
    # builder.add_node("rewrite_query", wyx_rewrite_query_node)
    builder.add_node("biw_query", biw_query_node)
    builder.add_node("dataset_filter", dataset_filter_node)
    builder.add_edge("dataset_filter", END)

    return builder


def build_graph_with_memory():
    """Build and return the agent workflow graph with memory."""
    # use persistent memory to save conversation history
    # TODO: be compatible with SQLite / PostgreSQL
    memory = MemorySaver()

    builder = _build_base_graph()
    return builder.compile(checkpointer=memory)


def build_graph():
    """Build and return the agent workflow graph without memory."""
    builder = _build_base_graph()
    return builder.compile()